
from PIL import Image
import os,json
from torch.utils.data import Dataset

def pil_loader(path):
    return Image.open(path).convert('RGB')

class FundusDataset(Dataset):
    def __init__(self, root, split='train', split_name='unsuperAD',transform=None):
        self.data = []
        self.root = root
        self.transform = transform
        with open(os.path.join(root,'annotations.json')) as f:
            self.data_dict=json.load(f)
        with open(os.path.join(root,'split',f'{split_name}.json')) as f:
            split_dict=json.load(f)
            if split not in split_dict:
                raise ValueError(f"Invalid split name: {split} in {split_name}.json")
            self.split_list=split_dict[split]

    def __getitem__(self, index):
        image_name=self.split_list[index]
        data=self.data_dict[image_name]
        image_path=os.path.join(self.root,data["resize_image_path"])
        image=pil_loader(image_path)
        if self.transform:
            image=self.transform(image)
        label=0 if data['stage']==0 else 1
        return image,{'label':label,'image_path':image_path}

    def __len__(self):
        return len(self.split_list)
